from langchain.agents import AgentType, initialize_agent
from langchain.chat_models import ChatOpenAI
from langchain.tools.base import BaseTool

# 定义一个自定义工具类
class Calculator(BaseTool):
    name = "Calculator"
    description = "运行加减乘除运算的表达式"

    def _run(self, query: str):
        """Use the tool."""
        # 在这里实现您的自定义函数逻辑
        result = self.custom_function(query)
        return result

    async def _arun(self, query: str):
        """Use the tool asynchronously."""
        raise NotImplementedError("This tool does not support async")

    def custom_function(self, expression: str):
        """Your custom function logic goes here."""
        # 示例：将输入文本转换为大写
        return eval(expression)

# 创建工具实例
calculator = Calculator()

tools = [calculator]

chat_model = ChatOpenAI()

# 初始化代理
agent = initialize_agent(tools, chat_model, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True)
response = agent.run("计算 10*10/2")

print(response)
